package vector

import (
	"context"
	"fmt"

	"github.com/qdrant/go-client/qdrant"
	"go-file-perception-model/internal/logger"
	"go.uber.org/zap"
)

// QdrantCollectionManager Qdrant集合管理器
type QdrantCollectionManager struct {
	client *QdrantClient
}

// NewQdrantCollectionManager 创建集合管理器
func NewQdrantCollectionManager(client *QdrantClient) *QdrantCollectionManager {
	return &QdrantCollectionManager{
		client: client,
	}
}

// CreateCollection 创建集合
func (m *QdrantCollectionManager) CreateCollection(ctx context.Context, collectionName string, vectorSize int) error {
	// 检查集合是否已存在
	exists, err := m.CollectionExists(ctx, collectionName)
	if err != nil {
		return fmt.Errorf("failed to check collection existence: %w", err)
	}

	if exists {
		logger.Info("Collection already exists", zap.String("collection", collectionName))
		return nil
	}

	// 创建集合
	createCollection := qdrant.CreateCollection{
		CollectionName: collectionName,
		VectorsConfig: &qdrant.VectorsConfig{
			Config: &qdrant.VectorsConfig_Params{
				Params: &qdrant.VectorParams{
					Size:     uint64(vectorSize),
					Distance: qdrant.Distance_Cosine, // 使用余弦相似度
				},
			},
		},
	}

	err = m.client.GetClient().CreateCollection(ctx, &createCollection)
	if err != nil {
		return fmt.Errorf("failed to create collection: %w", err)
	}

	logger.Info("Collection created successfully", zap.String("collection", collectionName))
	return nil
}

// CollectionExists 检查集合是否存在
func (m *QdrantCollectionManager) CollectionExists(ctx context.Context, collectionName string) (bool, error) {
	collections, err := m.client.GetClient().ListCollections(ctx)
	if err != nil {
		return false, fmt.Errorf("failed to list collections: %w", err)
	}

	for _, collection := range collections {
		if collection == collectionName {
			return true, nil
		}
	}

	return false, nil
}

// GetCollectionInfo 获取集合信息
func (m *QdrantCollectionManager) GetCollectionInfo(ctx context.Context, collectionName string) (*qdrant.CollectionInfo, error) {
	collectionInfo, err := m.client.GetClient().GetCollectionInfo(ctx, collectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to get collection info: %w", err)
	}

	return collectionInfo, nil
}

// DeleteCollection 删除集合
func (m *QdrantCollectionManager) DeleteCollection(ctx context.Context, collectionName string) error {
	err := m.client.GetClient().DeleteCollection(ctx, collectionName)
	if err != nil {
		return fmt.Errorf("failed to delete collection: %w", err)
	}

	logger.Info("Collection deleted successfully", zap.String("collection", collectionName))
	return nil
}

// ListCollections 列出所有集合
func (m *QdrantCollectionManager) ListCollections(ctx context.Context) ([]string, error) {
	collections, err := m.client.GetClient().ListCollections(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list collections: %w", err)
	}

	return collections, nil
}